# -*- coding: utf-8 -*-
"""
@Author: zyx
@Date: 2024/11/28 15:19
@FileName: data_processor.py
@Description: 文档数据处理
"""

import os
from loguru import logger
from time import time


def chroma_process(papers_path: str, tags_chroma_path: str, knowledge_chroma_path: str):
    from utils import text_loader_splitter, get_vectordb, build_tag_chromas

    """数据处理
    :param str papers_path: 论文所在目录
    :param str tags_chroma_path: 标签向量库位置
    :param str knowledge_chroma_path: 知识库位置
    """
    os.makedirs(tags_chroma_path, exist_ok=True)
    os.makedirs(knowledge_chroma_path, exist_ok=True)
    logger.info("****** 开始构建tag向量库 ******")
    build_tag_chromas(root_path=tags_chroma_path)
    logger.info("****** 构建tag向量库成功 ******")
    for coll in os.listdir(papers_path):
        # 按照论文目录建立知识库
        sub_path = os.path.join(papers_path, coll)
        tag_vectordb = get_vectordb(
            coll_path=os.path.join(tags_chroma_path, coll), coll=coll
        )
        knowledge_vectordb = get_vectordb(
            coll_path=os.path.join(knowledge_chroma_path, coll),
            coll=coll,
            is_reset=True,
        )
        for f in os.listdir(sub_path):  # 获取文档
            # 文档切分
            docs = text_loader_splitter(os.path.join(sub_path, f))
            for doc in docs:
                # 依据相似度为文档添加标签
                tag = tag_vectordb.similarity_search(doc.page_content, k=1)[0]
                doc.metadata["tag"] = tag.page_content
            # 文档向量入知识库
            knowledge_vectordb.add_documents(docs)
            logger.info(f"文档【{doc.metadata['source']}】向量处理完毕")


def lightrag_process(papers_txt_path: str, lightrag_path: str):
    """数据处理
    :param str papers_txt_path: 论文txt所在目录
    :param str lightrag_path: lightrag相关文件所在位置
    """
    from rag_wrapper import LightRAGWrapper

    os.makedirs(lightrag_path, exist_ok=True)
    lrw = LightRAGWrapper(working_dir=lightrag_path)
    for paper in os.listdir(papers_txt_path):
        with open(os.path.join(papers_txt_path, paper), "r", encoding="utf-8") as f:
            text = f.read()
        logger.info(f"********** start to handle paper {paper} **********")
        start = time()
        lrw.add_text(text)
        logger.info(f"========== success to handle paper {paper} ==========")
        logger.info(
            f"----->>>>> handle paper {paper} cost {time()-start:.2f}s <<<<<*****"
        )


if __name__ == "__main__":
    # chroma_process(
    #     papers_path="./datas/papers/",
    #     tags_chroma_path="./datas/tag_chroma/",
    #     knowledge_chroma_path="./datas/knowledge_chroma/",
    # )
    lightrag_process(
        papers_txt_path="./datas/papers-txt/", lightrag_path="./datas/papers_lightrag/"
    )
